var express = require('express');
var router = express.Router();
const { ollama, chain, saveMemory } = require('../utils/ollamaUtils');
const { HumanMessage } = require("@langchain/core/messages");

/* GET home page. */
router.get('/', function (req, res, next) {
  res.render('index', { title: 'lbc-ollama' });
});

// router.get('/model-list', (req, res) => {
//   ollama
//   // res.send()
// })

function setResponseHeaders(res) {
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');
  res.flushHeaders();
}

router.post('/chat', function (req, res, next) {
  let options = req.body;
  if (options.stream) {
    setResponseHeaders(res);
    ollama.stream(options.message).then(async (response) => {
      for await (const part of response) {
        res.write(part);
      }
      res.end();
    })
  } else {
    let response = ollama.invoke(HumanMessage(options.message));
    res.send(response);
  }
})

router.post('/chatPrompt', function (req, res, next) {
  let options = req.body;
  if (options.stream) {
    setResponseHeaders(res);
    let fullResponse = "";
    chain.stream({ input: options.message }).then(async (response) => {
      console.log(response, 1111111);
      for await (const part of response) {
        fullResponse += part.response;
        res.write(part.response);
      }
      await saveMemory({
        input: options.message,
        output: fullResponse
      })
      res.end();
    })
  } else {
    let response = chain.invoke({ input: options.message });
    res.send(response);
  }
})

router.get('/stopChat', function (req, res, next) {
  model.abort();
  res.end();
})





module.exports = router;


